import os
import pandas as pd
from langchain_text_splitters import RecursiveCharacterTextSplitter
from openai import OpenAI
from pymilvus import Collection, FieldSchema, CollectionSchema, DataType, connections
from langchain.schema import Document

os.environ["OPENAI_API_KEY"] = "sk-38b1a77d899b4e708287a296ceeb02e3"

def get_embeddings(input_text):
    client = OpenAI(
        api_key=os.getenv("OPENAI_API_KEY"),
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    )
    completion = client.embeddings.create(
        model="text-embedding-v3",
        input=input_text,
        encoding_format="float"
    )
    return completion.data[0].embedding

symp = pd.read_csv("SymptomTable.csv", encoding="utf-8")

#症状部分
symp_documents = [Document(page_content=text, metadata={"source": "SymptomTable"}) for text in symp["SymptomDescription"].dropna().tolist()]


#穴位名称数据库
text_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=200)
split_documents = text_splitter.split_documents(symp_documents)
embeddings = [get_embeddings(doc.page_content) for doc in split_documents]
connections.connect("default", host="localhost", port="19530")
collection_name = "symptom"
id_field = FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True)  #使用 auto_id 自动生成 ID
vector_field = FieldSchema(name="vector", dtype=DataType.FLOAT_VECTOR, dim=1024)  #embedding维度1024
text_field = FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=65535)
metadata_field = FieldSchema(name="metadata", dtype=DataType.JSON)
schema = CollectionSchema(fields=[id_field, vector_field, text_field, metadata_field], description="Symptom schema")
collection = Collection(name=collection_name, schema=schema)
texts = [doc.page_content for doc in split_documents]  #文档内容
metadatas = [doc.metadata for doc in split_documents]  #元数据
collection.insert([embeddings, texts, metadatas])
index_params = {
    "index_type": "IVF_FLAT",
    "metric_type": "L2",
    "params": {"nlist": 128}
}
collection.create_index(field_name="vector", index_params=index_params)
print(f"Index created for collection {collection_name}")
collection = Collection("symptom")


